import cv2
import torch
from facenet_pytorch import MTCNN, InceptionResnetV1
import numpy as np
import os

# 获取设备
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)

# mtcnn模型加载【设置网络参数，进行人脸检测】
mtcnn = MTCNN(min_face_size=12, thresholds=[0.2, 0.2, 0.3], keep_all=False, device=device)
# InceptionResnetV1模型加载【用于获取人脸特征向量】
resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device)

# 定义文件夹路径
input_folder = 'true'
output_folder = 'tz'

# 确保输出文件夹存在
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# 遍历文件夹中的所有图片文件
for filename in os.listdir(input_folder):
    if filename.endswith(('.png', '.jpg', '.jpeg')):
        img_path = os.path.join(input_folder, filename)
        img = cv2.imread(img_path)  # 读取图片
        if img is None:
            print(f"Image not found at path: {img_path}")
            continue

        face = mtcnn(img)  # 使用mtcnn检测人脸，返回人脸张量
        if face is not None:
            # 添加一个批次维度（在0维位置）
            face = face.unsqueeze(0).to(device)
        else:
            print(f"No faces detected in the image: {img_path}")
            continue

        # 确保输入张量是四维的
        if face.dim() == 4:
            with torch.no_grad():
                face_emb = resnet(face).detach().cpu()
        else:
            print(f"The input tensor must be 4-dimensional for image: {img_path}")
            continue

        # 保存特征向量为.npy文件
        feature_vector_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.npy")
        np.save(feature_vector_path, face_emb)
        print(f"特征向量已保存为: {feature_vector_path}")
